# -*- coding: utf-8 -*-
"""
测试rerank框架的基本功能
"""

import os
import sys
from typing import List
from loguru import logger

# 添加项目路径
sys.path.append(os.path.dirname(os.path.dirname(__file__)))

from questionretrieval.retrievers.rerank import (
    BGEReranker,
    EmbeddingRerankComparison,
    RerankOnlyComparison
)

# 配置loguru日志
logger.remove()  # 移除默认handler
logger.add(
    sys.stderr,
    format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
    level="INFO"
)


def test_individual_rerankers():
    """
    测试单个reranker模型
    """
    logger.info("开始测试单个reranker模型")
    
    # 测试数据
    query = "什么是机器学习？"
    documents = [
        "机器学习是人工智能的一个分支，它使计算机能够在没有明确编程的情况下学习。",
        "深度学习是机器学习的一个子集，使用多层神经网络来模拟人脑的工作方式。",
        "监督学习需要标记的训练数据，而无监督学习则不需要标签。",
        "梯度下降是一种优化算法，用于最小化损失函数。",
        "Python是一种流行的编程语言，广泛用于数据科学和机器学习。"
    ]
    
    # 测试BGE Reranker Base
    try:
        logger.info("测试BGE Reranker Base")
        bge_base = BGEReranker('bge_base', device='cpu')  # 使用CPU避免GPU内存问题
        bge_base.load_model()
        
        results = bge_base.rerank(query, documents, top_k=3)
        logger.info(f"BGE Base结果: {len(results)} 个文档")
        for i, (doc, score) in enumerate(results):
            logger.info(f"  {i+1}. 分数: {score:.4f}, 文档: {doc[:50]}...")
            
    except Exception as e:
        logger.error(f"BGE Reranker Base测试失败: {e}")
    



def test_embedding_rerank_comparison():
    """
    测试embedding+rerank对比
    """
    logger.info("开始测试embedding+rerank对比")
    
    # 测试数据
    queries = [
        "什么是机器学习？",
        "深度学习的原理是什么？"
    ]
    
    documents = [
        "机器学习是人工智能的一个分支，它使计算机能够在没有明确编程的情况下学习。",
        "深度学习是机器学习的一个子集，使用多层神经网络来模拟人脑的工作方式。",
        "监督学习需要标记的训练数据，而无监督学习则不需要标签。",
        "梯度下降是一种优化算法，用于最小化损失函数。",
        "Adam优化器结合了动量和自适应学习率的优点。",
        "反向传播算法通过计算梯度来更新神经网络的权重。",
        "过拟合发生在模型在训练数据上表现很好，但在测试数据上表现较差时。",
        "正则化技术如L1和L2可以帮助防止过拟合。",
        "交叉验证是评估模型性能的重要技术。",
        "特征工程对机器学习模型的性能有重要影响。"
    ]
    
    try:
        # 初始化对比实验（使用较小的参数进行测试）
        experiment = EmbeddingRerankComparison(
            embedding_model='BAAI/bge-m3',
            top_k_embedding=8,  # 减少候选数量
            top_k_rerank=3,     # 减少最终结果数量
            device='cpu'        # 使用CPU
        )
        
        # 加载embedding模型和文档
        logger.info("加载embedding模型...")
        experiment.load_embedding_model()
        
        logger.info("加载文档...")
        experiment.load_documents(documents)
        
        # 运行单个查询测试
        logger.info("运行单个查询测试...")
        result = experiment.run_single_comparison(queries[0])
        
        logger.info(f"单个查询结果包含 {len(result)} 个reranker的结果")
        for model_name, model_result in result.items():
            logger.info(f"  {model_name}: {len(model_result.final_ranking)} 个文档")
        
        # 运行批量测试（只测试前2个查询）
        logger.info("运行批量查询测试...")
        results = experiment.run_batch_comparison(queries[:2])
        
        logger.info(f"批量查询结果包含 {len(results)} 个查询的结果")
        
    except Exception as e:
        logger.error(f"Embedding+Rerank对比测试失败: {e}")
        import traceback
        traceback.print_exc()


def test_rerank_only_comparison():
    """
    测试rerank only对比
    """
    logger.info("开始测试rerank only对比")
    
    # 测试数据
    queries = [
        "什么是机器学习？",
        "深度学习的原理是什么？"
    ]
    
    documents = [
        "机器学习是人工智能的一个分支，它使计算机能够在没有明确编程的情况下学习。",
        "深度学习是机器学习的一个子集，使用多层神经网络来模拟人脑的工作方式。",
        "监督学习需要标记的训练数据，而无监督学习则不需要标签。",
        "梯度下降是一种优化算法，用于最小化损失函数。",
        "Adam优化器结合了动量和自适应学习率的优点。",
        "反向传播算法通过计算梯度来更新神经网络的权重。",
        "过拟合发生在模型在训练数据上表现很好，但在测试数据上表现较差时。",
        "正则化技术如L1和L2可以帮助防止过拟合。"
    ]
    
    try:
        # 初始化对比实验
        experiment = RerankOnlyComparison(
            top_k=3,
            device='cpu',
            enable_parallel=False  # 禁用并行以简化测试
        )
        
        # 加载模型（只加载可用的模型）
        logger.info("加载rerank模型...")
        available_models = []
        
        # 尝试加载所有模型
        try:
            experiment.load_all_models()
            available_models = list(experiment.rerankers.keys())
            logger.info(f"成功加载模型: {available_models}")
        except Exception as e:
            logger.warning(f"模型加载失败: {e}")
        

        
        if not available_models:
            logger.error("没有可用的rerank模型")
            return
        
        logger.info(f"可用模型: {available_models}")
        
        # 运行统一候选集对比
        logger.info("运行统一候选集对比...")
        results = experiment.run_uniform_candidates_comparison(queries[:1], documents)
        
        logger.info(f"统一候选集对比结果包含 {len(results)} 个查询的结果")
        for query_idx, query_result in enumerate(results):
            logger.info(f"  查询 {query_idx}: {len(query_result)} 个模型的结果")
            for model_name, model_result in query_result.items():
                logger.info(f"    模型 {model_name}: 处理时间 {model_result.processing_time:.3f}秒")
        
    except Exception as e:
        logger.error(f"Rerank Only对比测试失败: {e}")
        import traceback
        traceback.print_exc()


def main():
    """
    主测试函数
    """
    logger.info("开始rerank框架测试")
    
    # 检查环境
    logger.info(f"Python版本: {sys.version}")
    logger.info(f"工作目录: {os.getcwd()}")
    
    # 检查模型可用性
    logger.info("准备测试BGE模型")
    
    try:
        # 1. 测试单个reranker
        logger.info("=" * 50)
        test_individual_rerankers()
        
        # 2. 测试embedding+rerank对比
        logger.info("=" * 50)
        test_embedding_rerank_comparison()
        
        # 3. 测试rerank only对比
        logger.info("=" * 50)
        test_rerank_only_comparison()
        
        logger.info("=" * 50)
        logger.info("所有测试完成！")
        
    except Exception as e:
        logger.error(f"测试过程中发生错误: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()